import torch
import torch_geometric
import gzip
import pickle
from torch_geometric.data import HeteroData


class GraphNodeDataset(torch_geometric.data.Dataset):
    """
    Constructs graph dataset from Node features
    """

    def __init__(self, data_files):
        super().__init__(root=None, transform=None, pre_transform=None)
        self.data_files = data_files

    def len(self):
        return len(self.data_files)

    def get(self, index):
        """
        This method return a node bipartite graph sample as saved on the disk during data collection.
        """
        with gzip.open(self.data_files[index], 'rb') as f:
            sample = pickle.load(f)
        features, target = sample[0], sample[1]

        data = HeteroData()
        # add node feature information
        data['var'].x = torch.tensor(features.variable_features,dtype=torch.float)
        data['con'].x = torch.tensor(features.constraint_features,dtype=torch.float)
        data['met'].x = torch.tensor(features.metric_features,dtype=torch.float)
        # add edge index information
        data['var','v_c','con'].edge_index = torch.tensor(features.vc_edge_index,dtype=torch.int64)
        data['con','c_v','var'].edge_index = torch.tensor(features.cv_edge_index,dtype=torch.int64)
        data['var','v_m','met'].edge_index = torch.tensor(features.vm_edge_index,dtype=torch.int64)
        data['met','m_v','var'].edge_index = torch.tensor(features.mv_edge_index,dtype=torch.int64)
        # # No edge attribute
        # data['v_c'].edge_attr = features.vc_edge_attr
        # data['c_v'].edge_attr = features.cv_edge_attr
        # data['v_m'].edge_attr = features.vm_edge_attr
        # data['m_v'].edge_attr = features.mv_edge_attr

        # Add label for the graph classification task
        # Since the label is a graph-level attribute, we add it to the data object
        data['graph_label'] = torch.tensor(1.,dtype=torch.float) if target else torch.tensor(0.,dtype=torch.float)

        return data
